import pandas as pd
from sqlalchemy import create_engine
from urllib.parse import quote_plus

stations = [
    'Aotizhongxin',
    'Changping',
    'Dingling',
    'Dongsi',
    'Guanyuan',
    'Gucheng',
    'Huairou',
    'Nongzhanguan',
    'Shunyi',
    'Tiantan',
    'Wanliu',
    'Wanshouxigong'
]

engine = create_engine(f'mysql+pymysql://root:{quote_plus("123456")}@127.0.0.1:3306/air-quality')

# 定义你的SQL查询
sql_query = 'select * from t_prsa_data where station = "%s" order by `year`, `month`, `day`, `hour`'

for station in stations:
    df = pd.read_sql_query(sql_query % station, engine)

    df['datetime'] = pd.to_datetime(df[['year', 'month', 'day', 'hour']])
    df.set_index('datetime', inplace=True)

    df.drop(['year', 'month', 'day', 'hour', 'station'], axis=1, inplace=True)
    air_quality_day_df = df[['pm25', 'pm10', 'so2', 'no2', 'co']].resample('D').mean()

    # Calculate the daily maximum 1-hour average for O3
    air_quality_day_df['o3_max'] = df['o3'].resample('D').max()

    # Calculate the daily maximum 8-hour sliding average for O3
    # For this, we'll use a rolling window of 8 hours and then resample to daily maximum
    air_quality_day_df['o3_8_avg_max'] = df['o3'].rolling(window='8h').mean().resample('D').max()

    air_quality_day_df.reset_index(inplace=True)
    air_quality_day_df['year'] = air_quality_day_df['datetime'].dt.year
    air_quality_day_df['month'] = air_quality_day_df['datetime'].dt.month
    air_quality_day_df['day'] = air_quality_day_df['datetime'].dt.day
    air_quality_day_df.drop(['datetime'], axis=1, inplace=True)

    air_quality_day_df['station'] = station

    air_quality_day_df.to_sql(name='t_air_quality_day', con=engine, if_exists='append', index=False)
